#!/usr/bin/env python3
"""
数据下载脚本

支持从多个交易所下载K线数据，按配置文件指定的结构保存
"""

import sys
import os
import argparse
from datetime import datetime, timedelta
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from xcquant.core import config_manager, unified_logger
from xcquant.data.data_manager import DataManager
# 移除不存在的导入
# from xcquant.data.data_source import DataSource


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='数据下载工具')
    
    parser.add_argument('--config', '-c', 
                       default='config.yaml',
                       help='配置文件路径')
    
    parser.add_argument('--exchange', '-e',
                       choices=['okx', 'binance', 'all'],
                       default='all',
                       help='交易所选择')
    
    parser.add_argument('--symbol', '-s',
                       help='交易对 (如: BTC-USDT)')
    
    parser.add_argument('--symbols',
                       help='交易对列表文件路径')
    
    parser.add_argument('--interval', '-i',
                       default='1m',
                       help='时间间隔 (如: 1m, 5m, 1h, 1d)')
    
    parser.add_argument('--type', '-t',
                       choices=['spot', 'futures'],
                       default='spot',
                       help='交易类型')
    
    parser.add_argument('--days', '-d',
                       type=int,
                       default=30,
                       help='下载天数')
    
    parser.add_argument('--start-date',
                       help='开始日期 (YYYY-MM-DD)')
    
    parser.add_argument('--end-date',
                       help='结束日期 (YYYY-MM-DD)')
    
    parser.add_argument('--update', '-u',
                       action='store_true',
                       help='更新现有数据')
    
    parser.add_argument('--force', '-f',
                       action='store_true',
                       help='强制重新下载')
    
    parser.add_argument('--list-symbols',
                       action='store_true',
                       help='列出可用交易对')
    
    return parser.parse_args()


def load_symbols_from_file(file_path: str) -> list:
    """从文件加载交易对列表"""
    symbols = []
    try:
        with open(file_path, 'r') as f:
            for line in f:
                symbol = line.strip()
                if symbol and not symbol.startswith('#'):
                    symbols.append(symbol)
    except Exception as e:
        print(f"加载交易对文件失败: {e}")
    
    return symbols


def get_default_symbols(exchange: str, trade_type: str) -> list:
    """获取默认交易对列表"""
    if trade_type == 'spot':
        return [
            'BTC-USDT', 'ETH-USDT', 'BNB-USDT', 'ADA-USDT', 'XRP-USDT',
            'SOL-USDT', 'DOT-USDT', 'DOGE-USDT', 'AVAX-USDT', 'MATIC-USDT'
        ]
    else:  # futures
        return [
            'BTC-USDT-SWAP', 'ETH-USDT-SWAP', 'BNB-USDT-SWAP', 
            'ADA-USDT-SWAP', 'XRP-USDT-SWAP'
        ]


def create_exchange_client(exchange_name: str, config: dict):
    """创建交易所客户端"""
    exchange_config = config.get('exchanges', {}).get(exchange_name, {})
    
    if exchange_name == 'okx':
        return OKXExchange(exchange_config)
    elif exchange_name == 'binance':
        return BinanceExchange(exchange_config)
    else:
        raise ValueError(f"不支持的交易所: {exchange_name}")


def download_data_for_exchange(data_manager: DataManager, exchange_name: str, 
                             symbols: list, interval: str, trade_type: str,
                             start_time: datetime, end_time: datetime,
                             force: bool, logger):
    """为指定交易所下载数据"""
    logger.info(f"开始下载 {exchange_name} 数据...")
    
    success_count = 0
    total_count = len(symbols)
    
    for i, symbol in enumerate(symbols, 1):
        try:
            logger.info(f"[{i}/{total_count}] 下载 {symbol} {interval} 数据...")
            
            # 检查本地数据
            if not force:
                local_data = data_manager.load_data(exchange_name, symbol, interval, 
                                                  trade_type, start_time, end_time)
                if local_data is not None and not local_data.empty:
                    logger.info(f"本地数据已存在，跳过: {symbol}")
                    success_count += 1
                    continue
            
            # 下载数据
            success = data_manager.download_and_save_data(
                exchange_name, symbol, interval, trade_type, start_time, end_time
            )
            
            if success:
                success_count += 1
                logger.info(f"下载成功: {symbol}")
            else:
                logger.error(f"下载失败: {symbol}")
            
            # 避免请求过于频繁
            import time
            time.sleep(0.1)
            
        except Exception as e:
            logger.error(f"下载 {symbol} 数据时出错: {e}")
    
    logger.info(f"{exchange_name} 下载完成: {success_count}/{total_count}")
    return success_count


def update_data_for_exchange(data_manager: DataManager, exchange_name: str,
                           symbols: list, interval: str, trade_type: str, logger):
    """更新指定交易所的数据"""
    logger.info(f"开始更新 {exchange_name} 数据...")
    
    success_count = 0
    total_count = len(symbols)
    
    for i, symbol in enumerate(symbols, 1):
        try:
            logger.info(f"[{i}/{total_count}] 更新 {symbol} {interval} 数据...")
            
            success = data_manager.update_data(exchange_name, symbol, interval, trade_type)
            
            if success:
                success_count += 1
                logger.info(f"更新成功: {symbol}")
            else:
                logger.error(f"更新失败: {symbol}")
            
            # 避免请求过于频繁
            import time
            time.sleep(0.1)
            
        except Exception as e:
            logger.error(f"更新 {symbol} 数据时出错: {e}")
    
    logger.info(f"{exchange_name} 更新完成: {success_count}/{total_count}")
    return success_count


def list_available_symbols(exchange_name: str, config: dict, logger):
    """列出可用的交易对"""
    try:
        exchange_client = create_exchange_client(exchange_name, config)
        
        # 这里需要根据实际的交易所API实现
        # 暂时返回默认列表
        spot_symbols = get_default_symbols(exchange_name, 'spot')
        futures_symbols = get_default_symbols(exchange_name, 'futures')
        
        print(f"\n{exchange_name.upper()} 可用交易对:")
        print("现货:")
        for symbol in spot_symbols:
            print(f"  {symbol}")
        
        print("合约:")
        for symbol in futures_symbols:
            print(f"  {symbol}")
        
    except Exception as e:
        logger.error(f"获取 {exchange_name} 交易对列表失败: {e}")


def main():
    """主函数"""
    args = parse_arguments()
    
    # 加载配置
    try:
        config_manager = ConfigManager(args.config)
        config = config_manager.get_config()
    except Exception as e:
        print(f"加载配置文件失败: {e}")
        sys.exit(1)
    
    # 初始化日志
    logger = get_logger("data_downloader")
    
    # 列出交易对
    if args.list_symbols:
        exchanges = ['okx', 'binance'] if args.exchange == 'all' else [args.exchange]
        for exchange in exchanges:
            list_available_symbols(exchange, config, logger)
        return
    
    # 初始化数据管理器
    try:
        data_config = config.get('data', {})
        data_manager = DataManager(data_config)
        
        # 注册交易所
        exchanges_config = config.get('exchanges', {})
        for exchange_name in ['okx', 'binance']:
            if exchange_name in exchanges_config:
                exchange_client = create_exchange_client(exchange_name, config)
                data_manager.register_exchange(exchange_name, exchange_client)
        
    except Exception as e:
        logger.error(f"初始化数据管理器失败: {e}")
        sys.exit(1)
    
    # 确定交易所列表
    if args.exchange == 'all':
        exchanges = ['okx', 'binance']
    else:
        exchanges = [args.exchange]
    
    # 确定交易对列表
    symbols = []
    if args.symbol:
        symbols = [args.symbol]
    elif args.symbols:
        symbols = load_symbols_from_file(args.symbols)
    else:
        # 使用默认交易对
        symbols = get_default_symbols(args.exchange, args.type)
    
    if not symbols:
        logger.error("没有指定交易对")
        sys.exit(1)
    
    # 确定时间范围
    if args.start_date and args.end_date:
        start_time = datetime.strptime(args.start_date, '%Y-%m-%d')
        end_time = datetime.strptime(args.end_date, '%Y-%m-%d')
    else:
        end_time = datetime.now()
        start_time = end_time - timedelta(days=args.days)
    
    logger.info(f"下载参数:")
    logger.info(f"  交易所: {exchanges}")
    logger.info(f"  交易对: {symbols}")
    logger.info(f"  时间间隔: {args.interval}")
    logger.info(f"  交易类型: {args.type}")
    logger.info(f"  时间范围: {start_time} - {end_time}")
    logger.info(f"  更新模式: {args.update}")
    logger.info(f"  强制下载: {args.force}")
    
    # 执行下载或更新
    total_success = 0
    total_count = 0
    
    for exchange_name in exchanges:
        if exchange_name not in data_manager._fetchers:
            logger.warning(f"跳过未注册的交易所: {exchange_name}")
            continue
        
        if args.update:
            success = update_data_for_exchange(
                data_manager, exchange_name, symbols, 
                args.interval, args.type, logger
            )
        else:
            success = download_data_for_exchange(
                data_manager, exchange_name, symbols, args.interval, args.type,
                start_time, end_time, args.force, logger
            )
        
        total_success += success
        total_count += len(symbols)
    
    # 输出统计信息
    logger.info(f"\n下载完成统计:")
    logger.info(f"  总成功: {total_success}")
    logger.info(f"  总数量: {total_count}")
    logger.info(f"  成功率: {total_success/total_count*100:.1f}%" if total_count > 0 else "  成功率: 0%")
    
    # 显示数据存储信息
    available_data = data_manager.get_available_data()
    if available_data:
        logger.info(f"\n可用数据:")
        for exchange, types in available_data.items():
            logger.info(f"  {exchange}:")
            for trade_type, symbols_data in types.items():
                logger.info(f"    {trade_type}: {len(symbols_data)} 个交易对")


if __name__ == '__main__':
    main()